from typing import List, Dict, Optional
from dataclasses import dataclass, field
from datetime import datetime

@dataclass
class DialogueContext:
    """对话上下文类"""
    history: List[Dict] = field(default_factory=list)
    last_query: Optional[str] = None
    last_result: Optional[Dict] = None
    mentioned_tables: List[str] = field(default_factory=list)
    current_focus: Optional[str] = None

class ContextManager:
    """负责管理对话上下文"""

    def __init__(self):
        self.context = DialogueContext()

    def update_context(self,
                       question: str,
                       query: str,
                       results: Dict,
                       tables: List[str]):
        """更新对话上下文"""
        self.context.history.append({
            'timestamp': datetime.now().isoformat(),
            'question': question,
            'query': query,
            'results': results
        })
        self.context.last_query = query
        self.context.last_result = results
        self.context.mentioned_tables = list(set(self.context.mentioned_tables) | set(tables))

    def get_context(self) -> DialogueContext:
        """获取当前上下文"""
        return self.context
